from scipy import ndimage
import cv2
import numpy as np
from configs.config import Config
import logging

class SensorProcessor:
    """传感器数据处理类，负责图像和激光雷达数据的处理"""

    def __init__(self):
        config = Config()
        self.img_width = config.get('airsim_gym3.img_width', 800)
        self.img_height = config.get('airsim_gym3.img_height', 600)
        self.lidar_points = config.get('airsim_gym3.lidar_points', 1000)
        self.depth_scale = config.get('Vision_process.depth_scale', 1000.0)
        self.max_valid_depth = config.get('Vision_process.max_valid_depth', 100.0)
        self.num_distance_sensors = config.get('Vision_process.num_distance_sensors', 16)
        # self.lidar_nun_sectors = config.get('Vision_process.lidar_num_sectors', 16)

        self.sensor_angles = np.linspace(0, 360, self.num_distance_sensors, endpoint=False)

    def process_rgb_image(self, response):
        """处理RGB图像数据"""
        if response.width == 0 or response.height == 0:
            raise RuntimeError(f"无效的RGB图像尺寸: {response.width}x{response.height}")

        img1d = np.frombuffer(response.image_data_uint8, dtype=np.uint8)
        img_rgb = img1d.reshape(response.height, response.width, 3)

        # 调整图像尺寸
        if img_rgb.shape != (self.img_height, self.img_width, 3):
            img_rgb = cv2.resize(img_rgb, (self.img_width, self.img_height))

        return img_rgb

    def process_depth_image(self, response):
        """处理深度图像数据"""
        if response.image_data_float is None:
            return np.zeros((self.img_height, self.img_width), dtype=np.uint16)

        try:
            # 转换为浮点数深度图
            depth_float = np.array(response.image_data_float, dtype=np.float32).reshape(
                response.height, response.width
            )

            # 处理异常值
            depth_float[np.isinf(depth_float)] = self.max_valid_depth
            depth_float[np.isnan(depth_float)] = 0.0
            depth_float[depth_float > self.max_valid_depth] = self.max_valid_depth

            # 应用中值滤波
            depth_filtered = ndimage.median_filter(depth_float, size=3)

            # 归一化深度值到 [0, 1] 范围
            normalized_depth = depth_filtered / self.max_valid_depth

            # 转换为 uint16 类型，最大值映射到 65535
            depth_uint16 = (normalized_depth * 65535).astype(np.uint16)

            # 调整尺寸
            if depth_uint16.shape != (self.img_height, self.img_width):
                depth_uint16 = cv2.resize(depth_uint16, (self.img_width, self.img_height),
                                          interpolation=cv2.INTER_NEAREST)

            return depth_uint16

        except Exception as e:
            logging.error(f"深度图处理错误: {e}")
            return np.zeros((self.img_height, self.img_width), dtype=np.uint16)

    def process_lidar_data(self, points):
        """处理激光雷达点云数据"""
        if len(points) == 0:
            return np.zeros((self.lidar_points, 3), dtype=np.float32)

        # 采样至固定点数
        if len(points) > self.lidar_points:
            indices = np.random.choice(len(points), self.lidar_points, replace=False)
            points = points[indices]
        else:
            # 补零填充
            padding = np.zeros((self.lidar_points - len(points), 3), dtype=np.float32)
            points = np.vstack([points, padding])

        return points

    def process_distance_sensor_data(self, sensor_data_list):
        """
        处理多个距离传感器的数据

        参数:
            sensor_data_list: 包含16个传感器数据的列表，每个元素是一个字典，包含:
                - sensor_id: 传感器ID
                - distance: 检测到的距离
                - position: 传感器位置 [x, y, z]
                - orientation: 传感器朝向 [roll, pitch, yaw]
        """
        if not sensor_data_list or len(sensor_data_list) != self.num_distance_sensors:
            logging.warning(f"期望{self.num_distance_sensors}个传感器数据，实际收到{len(sensor_data_list)}个")
            # 返回默认数据（全零）
            return np.zeros((self.num_distance_sensors, 3), dtype=np.float32)

        # 创建输出数组 [传感器ID, 角度, 距离]
        processed_data = np.zeros((self.num_distance_sensors, 3), dtype=np.float32)

        for sensor_data in sensor_data_list:
            sensor_id = sensor_data.get('sensor_id', -1)
            if sensor_id < 0 or sensor_id >= self.num_distance_sensors:
                continue

            # 获取距离值，无效值设为最大距离
            distance = sensor_data.get('distance', self.max_valid_depth)
            if distance <= 0 or distance > self.max_valid_depth:
                distance = self.max_valid_depth

            # 获取传感器朝向（假设使用yaw角）
            orientation = sensor_data.get('orientation', [0, 0, 0])
            yaw_angle = orientation[2]  # yaw角（绕Z轴旋转）

            # 存储 [传感器ID, 角度, 距离]
            processed_data[sensor_id] = [sensor_id, yaw_angle, distance]

        return processed_data
